import WebSocket from 'ws';

export function setupWebSocket(server) {
  const wss = new WebSocket.Server({ server });
  
  // 存储所有连接的客户端
  const clients = new Set();

  wss.on('connection', (ws) => {
    clients.add(ws);
    
    ws.on('close', () => {
      clients.delete(ws);
    });
  });

  // 广播消息给所有连接的客户端
  return function broadcast(message) {
    clients.forEach(client => {
      if (client.readyState === WebSocket.OPEN) {
        client.send(JSON.stringify(message));
      }
    });
  };
} 